//
// Created by xinyang on 2021/3/7.
//
// Modified by Harry-hhj on 2021/05/04

#include "TRTModule.hpp"

#include <thread>
#include <chrono>
//#include <common.hpp>
#include <fmt/format.h>
#include <fmt/color.h>
#include <opencv2/opencv.hpp>
#include <string>
#include <cmath>

using namespace std::chrono;

int main() {
    TRTModule model("/home/wmj/Desktop/sjt_test/model-opt-4.onnx");
    cv::VideoCapture cam("/home/wmj/Desktop/sjt_test/test.avi");
    cv::Mat img;

    const cv::Scalar colors[3] = {{255, 0, 0}, {0, 0, 255}, {0, 255, 0}};

    while (true) {
	    cam >> img;
        //cv::imshow("show first",img);
            auto detections = model(img);
            if(1){
                cv::Mat im2show = img.clone();
                for (const auto &b: detections) {
                    cv::line(im2show, b.pts[0], b.pts[1], colors[2], 2);
                    cv::line(im2show, b.pts[1], b.pts[2], colors[2], 2);
                    cv::line(im2show, b.pts[2], b.pts[3], colors[2], 2);
                    cv::line(im2show, b.pts[3], b.pts[0], colors[2], 2);
                    cv::putText(im2show, std::to_string(b.tag_id), b.pts[0], cv::FONT_HERSHEY_SIMPLEX, 1, colors[b.color_id]);
                }
                cv::imshow("img_show",im2show);
                cv::waitKey(10);
            }
    }
    return 0;
}

